import os
import cv2
import numpy as np


# 中值滤波
def median_image(img):
    if img is None or img.size == 0:
        raise ValueError("提供的图像为空，无法进行中值滤波处理")
    img_median = cv2.medianBlur(img, 5)
    return img_median


# 均值滤波
def mean_filter(image, kernel_size):
    if image is None or image.size == 0:
        raise ValueError("提供的图像为空，无法进行均值滤波处理")
        # 创建一个平均滤波器的卷积核
    kernel = np.ones((kernel_size, kernel_size), np.float32) / (kernel_size * kernel_size)

    # 使用滤波器进行图像滤波
    filtered_image = cv2.filter2D(image, -1, kernel)

    return filtered_image


# Canny边缘提取
def get_canny(img, low, high, sizes):
    if img is None or img.size == 0:
        raise ValueError("提供的图像为空，无法进行Canny边缘提取")
    edge = cv2.Canny(img, low, high, sizes)
    return edge


# 判断连通区域并且选择连通区域面积最大
def find_connected_components(img):
    if img is None or img.size == 0:
        raise ValueError("提供的图像为空，无法进行连通区域分析")
        # 查找连通区域
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img, connectivity=8)

    max_area_index = np.argmax(stats[1:, cv2.CC_STAT_AREA]) + 1

    areas = stats[max_area_index, cv2.CC_STAT_AREA]

    masks = np.zeros_like(img)
    masks[labels == max_area_index] = 255

    return masks, areas


def solve(img):
    # 使用中值滤波
    median_img = median_image(img)
    # 边缘提取
    edge_img = get_canny(median_img, 100, 200, 5)
    # 均值滤波
    mean_img = mean_filter(edge_img, 9)
    # 最大连通区域
    coned_img, s = find_connected_components(mean_img)

    vector_img = [img, median_img, edge_img, mean_img, coned_img]

    return vector_img, s


def assess(path1, path2):
    if not os.path.exists(path1):
        print("打击前图片路径不存在")
        return -1
    if not os.path.exists(path2):
        print("打击后图片路径不存在")
        return -2

    img1 = cv2.imdecode(np.fromfile(path1, dtype=np.uint8), -1)
    img2 = cv2.imdecode(np.fromfile(path2, dtype=np.uint8), -1)

    if img1 is None or img1.size == 0:
        print("打击前图像读取失败或为空")
        return -3
    if img2 is None or img2.size == 0:
        print("打击后图像读取失败或为空")
        return -4

    vector_img1, s1 = solve(img1)
    vector_img2, s2 = solve(img2)

    cnt = np.abs(s1 - s2)

    if cnt < 0.01 * s1:
        return '无损伤'
    elif cnt < 0.1 * s1:
        return '轻度损伤'
    else:
        return '重度损伤'


if __name__ == "__main__":
    # 打击前图片未毁伤
    path1 = "./pic/before.png"
    # 打击后图像
    path2 = "./pic/after.png"
    ret = assess(path1, path2)
    print(ret)
